////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//                                                                                                                            //
//                                                         WNProject                                                          //
//                                                                                                                            //
//            This file is distributed under the BSD 2-Clause open source license. See License.txt for details.               //
//                                                                                                                            //
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

#include "WNSystemInfo/inc/WNProcessorInfo.h"
#include "WNConcurrency/inc/WNExecuteOnce.h"
#include "WNMemory/inc/WNMemory.h"

#ifdef _WN_WINDOWS
    #include <windows.h>

    typedef WN_UINT32 (WINAPI *LPFN_GLPI)(PSYSTEM_LOGICAL_PROCESSOR_INFORMATION, PDWORD);
#elif defined _WN_LINUX
    #include <unistd.h>
    #include <stdio.h>
#endif


#ifdef _WN_X86
    #include "WNCore/inc/WNMacros.h"

    #ifdef _WN_MSVC
        #include <intrin.h>

        #if _WN_MSVC >= 160040219
            #include <immintrin.h>

            #define WNXGETBV(_x) _xgetbv(_x)
        #else
           #ifdef _WN_32_BIT
                #define xgetbv __asm _emit 0x0F __asm _emit 0x01 __asm _emit 0xD0

                __WN_INLINE WN_UINT64 WNXGETBV(WN_UINT32 _x) {
                    WN_UINT32 result1 = 0;
                    WN_UINT32 result2 = 0;

                    __asm {
                        mov ecx, _x
                        xgetbv
                        mov result1, eax
                        mov result2, edx
                    }

                    return((static_cast<WN_UINT64>(result2) << 32) | result1);
                }
            #elif defined _WN_64_BIT
                extern "C" WN_UINT64 WNXGETBV(WN_UINT32 _x);
            #endif

            #define _XCR_XFEATURE_ENABLED_MASK    0
        #endif

        #ifdef _WN_32_BIT
            __WN_INLINE WN_BOOL WNCPUIDSupported() {
                WN_UINT32 result = 0;

                __asm {
                    pushfd
                    pop eax
                    mov ecx, eax
                    xor eax, 0x200000
                    push eax
                    popfd
                    pushfd
                    pop eax
                    xor eax, ecx
                    mov result, eax
                    push ecx
                    popfd
                }

                return(result != 0);
            }
        #elif defined _WN_64_BIT
            #define WNCPUIDSupported() WN_TRUE
        #endif

        #define WNCPUID(_a, _b) __cpuid(_a, _b)

        #if _WN_MSVC < 150030729
            #ifdef _WN_32_BIT
                __WN_INLINE WN_VOID WNCPUIDEX(WN_INT32 _a[4], WN_INT32 _b, WN_INT32 _c) {
                    __asm {
                        // Load function type and ecx value
                        mov eax, _b
                        mov ecx, _c
                        cpuid
                        // Load info array address and store values into info array
                        mov edi, _a
                        mov [edi], eax
                        mov 4[edi], ebx
                        mov 8[edi], ecx
                        mov 12[edi], edx
                    }
                }
            #elif defined _WN_64_BIT
                extern "C" WN_VOID WNCPUIDEX(WN_INT32 _a[4], WN_INT32 _b, WN_INT32 _c);
            #endif
        #else
            #define WNCPUIDEX(_a, _b, _c) __cpuidex(_a, _b, _c)
        #endif
    #elif defined _WN_GNUC
        #ifdef _WN_32_BIT
            __WN_INLINE WN_BOOL WNCPUIDSupported() {
                WN_UINT32 result = 0;

                __asm __volatile (
                    "pushfl\n"
                    "pop %%eax\n"
                    "mov %%eax, %%ecx\n"
                    "xor $0x200000, %%eax\n"
                    "push %%eax\n"
                    "popfl\n"
                    "pushfl\n"
                    "pop %%eax\n"
                    "xor %%ecx, %%eax\n"
                    "mov %%eax, %0\n"
                    "push %%ecx\n"
                    "popfl\n" :
                    "=m"(result) : :
                    "eax",
                    "ecx",
                    "memory"
                );

                return(result != 0);
            }
        #elif defined _WN_64_BIT
            #define WNCPUIDSupported() WN_TRUE
        #endif

        #define _XCR_XFEATURE_ENABLED_MASK 0

        __WN_INLINE WN_UINT64 WNXGETBV(WN_UINT32 _x) {
            WN_UINT32 result1 = 0;
            WN_UINT32 result2 = 0;

            __asm__ __volatile__ (
                "xgetbv" :
                // Output
                "=a" (result1),
                "=d" (result2) :
                // Input
                "c" (_x)
            );

            return((static_cast<WN_UINT64>(result2) << 32) | result1);
        }

        __WN_INLINE WN_VOID WNCPUID(WN_INT32 _a[4], WN_INT32 _b) {
            __asm__ __volatile__ (
                "cpuid" :
                // Output
                "=a" (_a[0]),
                "=b" (_a[1]),
                "=c" (_a[2]),
                "=d" (_a[3]) :
                // Input
                "a" (_b)
            );
        }

        __WN_INLINE WN_VOID WNCPUIDEX(WN_INT32 _a[4], WN_INT32 _b, WN_INT32 _c) {
            __asm__ __volatile__ (
                "cpuid" :
                // Output
                "=a" (_a[0]),
                "=b" (_a[1]),
                "=c" (_a[2]),
                "=d" (_a[3]) :
                // Input
                "a" (_b),
                "c" (_c)
            );
        }
    #else
        #define WNCPUID(_a, _b)
        #define WNCPUIDEX(_a, _b, _c)
    #endif
#endif

static WN_BYTE g_Endianness = 0;
static WN_UINT8 g_PhysicalCoreCount = 1;
static WN_UINT8 g_LogicalThreadCount = 1;
static WN_UINT32 g_Features = 0;

WN_EXECUTE_ONCE_FUNCTION(WNProcessorInfo_Init) {
    // Determine endianness
    #ifdef _WN_ENDIAN_LITTLE
        g_Endianness = WN_PROCESSOR_ENDIAN_LITTLE;
    #elif defined _WN_ENDIAN_BIG
        g_Endianness = WN_PROCESSOR_ENDIAN_BIG;
    #else
        const WN_UINT32 value = 1;

        g_Endianness = *reinterpret_cast<const WN_BYTE*>(&value) == 1 ? WN_PROCESSOR_ENDIAN_LITTLE : WN_PROCESSOR_ENDIAN_BIG;
    #endif

    // Determine number of physical cores and logical threads
    #ifdef _WN_WINDOWS
        const LPFN_GLPI getLogicalProcessorInfo = reinterpret_cast<LPFN_GLPI>(GetProcAddress(GetModuleHandle(TEXT("kernel32")),
                                                                                             "GetLogicalProcessorInformation"));

        if (getLogicalProcessorInfo == WN_NULL) {
            SYSTEM_INFO sysInfo;

            GetSystemInfo(&sysInfo);

            if (sysInfo.dwNumberOfProcessors > 0) {
                g_LogicalThreadCount = static_cast<WN_UINT8>(sysInfo.dwNumberOfProcessors);
            }

            g_PhysicalCoreCount = g_LogicalThreadCount;

            #ifdef _WN_X86
                WN_INT32 info[4] = {0};

                WNCPUID(info, 1);

                if (WN_BIT_TEST(info[3], 28)) {
                    g_PhysicalCoreCount >>= 1;
                }
            #endif
        } else {
            PSYSTEM_LOGICAL_PROCESSOR_INFORMATION buffer = WN_NULL;

            DWORD returnLength = 0;

            while (WN_TRUE) {
                const WN_UINT32 returnCode = getLogicalProcessorInfo(buffer, &returnLength);

                if (returnCode == 0) {
                    if (GetLastError() == ERROR_INSUFFICIENT_BUFFER) {
                        if (buffer != WN_NULL) {
                            WNFree(buffer);
                        }

                        buffer = static_cast<PSYSTEM_LOGICAL_PROCESSOR_INFORMATION>(WNMalloc(returnLength));

                        if (buffer == WN_NULL) {
                            break;
                        }
                    } else {
                        break;
                    }
                } else {
                    break;
                }
            }

            if (buffer != WN_NULL) {
                g_PhysicalCoreCount = 0;
                g_LogicalThreadCount = 0;

                PSYSTEM_LOGICAL_PROCESSOR_INFORMATION ptr = buffer;

                WN_UINT32 byteOffset = 0;

                while (byteOffset + sizeof(SYSTEM_LOGICAL_PROCESSOR_INFORMATION) <= returnLength) {
                    if (ptr->Relationship == RelationProcessorCore) {
                        g_PhysicalCoreCount++;

                        const WN_SIZE_T shift = sizeof(WN_SIZE_T) * 8 - 1;
                        WN_SIZE_T bitTest = static_cast<WN_SIZE_T>(1) << shift;
                        WN_UINT8 bitSetCount = 0;

                        for (WN_SIZE_T i = 0; i <= shift; ++i) {
                            if ((ptr->ProcessorMask & bitTest) != 0) {
                                bitSetCount++;
                            }

                            bitTest >>= 1;
                        }

                        g_LogicalThreadCount += bitSetCount;
                    }

                    byteOffset += sizeof(SYSTEM_LOGICAL_PROCESSOR_INFORMATION);
                    ptr++;
                }

                WNFree(buffer);
            }
        }
    #elif defined _WN_LINUX
        // parse /proc/cpuinfo
    #endif

    // Determine available extended features
    #ifdef _WN_X86
        if (WNCPUIDSupported()) {
            WN_INT32 info[4] = {0};

            WNCPUID(info, 1);

            g_Features |= WN_BIT_TEST(info[3], 0) ? WN_PROCESSOR_FEATURE_FPU : 0;
            g_Features |= WN_BIT_TEST(info[3], 8) ? WN_PROCESSOR_FEATURE_X86_CMPXCHG8B : 0;
            g_Features |= WN_BIT_TEST(info[2], 13) ? WN_PROCESSOR_FEATURE_X86_CMPXCHG16B : 0;
            g_Features |= WN_BIT_TEST(info[3], 23) ? WN_PROCESSOR_FEATURE_X86_MMX : 0;
            g_Features |= WN_BIT_TEST(info[3], 25) ? WN_PROCESSOR_FEATURE_X86_SSE : 0;
            g_Features |= WN_BIT_TEST(info[3], 26) ? WN_PROCESSOR_FEATURE_X86_SSE2 : 0;
            g_Features |= WN_BIT_TEST(info[2], 0) ? WN_PROCESSOR_FEATURE_X86_SSE3 : 0;
            g_Features |= WN_BIT_TEST(info[2], 9) ? WN_PROCESSOR_FEATURE_X86_SSSE3 : 0;
            g_Features |= WN_BIT_TEST(info[2], 19) ? WN_PROCESSOR_FEATURE_X86_SSE4_1 : 0;
            g_Features |= WN_BIT_TEST(info[2], 20) ? WN_PROCESSOR_FEATURE_X86_SSE4_2 : 0;

            WNCPUID(info, 0x80000000);

            if (info[0] >= 0x80000001) {
                WNCPUID(info, 0x80000001);

                g_Features |= WN_BIT_TEST(info[3], 23) ? WN_PROCESSOR_FEATURE_X86_MMX_EXT : 0;
                g_Features |= WN_BIT_TEST(info[3], 31) ? WN_PROCESSOR_FEATURE_X86_3DNOW : 0;
                g_Features |= WN_BIT_TEST(info[3], 30) ? WN_PROCESSOR_FEATURE_X86_3DNOW_EXT : 0;
                g_Features |= WN_BIT_TEST(info[2], 6) ? WN_PROCESSOR_FEATURE_X86_SSE4A : 0;
            }

            WNCPUIDEX(info, 1, 1);

            const WN_BOOL OSXSAVESupported = WN_BIT_TEST(info[2], 27);
            const WN_BOOL AVXSupported = WN_BIT_TEST(info[2], 28);

            if (OSXSAVESupported && AVXSupported) {
                const WN_UINT64 XCRFeatureMask = WNXGETBV(_XCR_XFEATURE_ENABLED_MASK);

                g_Features |= (XCRFeatureMask & 0x6) == 1 ? WN_PROCESSOR_FEATURE_X86_AVX : 0;
            }
        }
    #endif

    if (g_PhysicalCoreCount < 1) {
        g_PhysicalCoreCount = 1;
    }

    if (g_LogicalThreadCount < 1) {
        g_LogicalThreadCount = 1;
    }
}

static WN_EXECUTE_ONCE_T g_Once = WN_EXECUTE_ONCE_INIT;

extern "C" {
    WN_UINT32 WNProcessorInfo_Features() {
        WNExecuteOnce(&g_Once, &WNProcessorInfo_Init);

        return(g_Features);
    }

    WN_BYTE WNProcessorInfo_Endianness() {
        WNExecuteOnce(&g_Once, &WNProcessorInfo_Init);

        return(g_Endianness);
    }

    WN_UINT8 WNProcessorInfo_PhysicalCoreCount() {
        WNExecuteOnce(&g_Once, &WNProcessorInfo_Init);

        return(g_PhysicalCoreCount);
    }

    WN_UINT8 WNProcessorInfo_LogicalThreadCount() {
        WNExecuteOnce(&g_Once, &WNProcessorInfo_Init);

        return(g_LogicalThreadCount);
    }
}